package com.shuigu.ai.langgraph4j.config;

import com.shuigu.ai.langgraph4j.CodeGenConcurrentWorkflow;
import com.shuigu.ai.langgraph4j.CodeGenWorkflow;
import org.bsc.langgraph4j.CompileConfig;
import org.bsc.langgraph4j.CompiledGraph;
import org.bsc.langgraph4j.StateGraph;
import org.bsc.langgraph4j.checkpoint.MemorySaver;
import org.bsc.langgraph4j.prebuilt.MessagesState;
import org.bsc.langgraph4j.state.AgentState;
import org.bsc.langgraph4j.studio.LangGraphStudioServer;
import org.bsc.langgraph4j.studio.springboot.LangGraphStudioConfig;
import org.springframework.context.annotation.Configuration;

import java.util.Map;

@Configuration
public class LangGraphStudioSampleConfig extends LangGraphStudioConfig {

    @Override
    public Map<String, LangGraphStudioServer.Instance> instanceMap() {

        CompiledGraph<MessagesState<String>> workflow = new CodeGenConcurrentWorkflow().createWorkflow();

        StateGraph<MessagesState<String>> stateGraph = workflow.stateGraph;

        // define your workflow

        var saver = new MemorySaver();

        var instance = LangGraphStudioServer.Instance.builder()
                .title("LangGraph Studio")
                .graph(stateGraph)
                .compileConfig(CompileConfig.builder()
                        .checkpointSaver(saver)
                        .build())
                .addInputStringArg("input")
                .build();

        return Map.of("default", instance);
    }
}
